#!/usr/bin/env python3
# Author: Armit
# Create Time: 2023/01/06 

from pathlib import Path
import gzip

from Bio.PDB import PDBParser
from Bio.PDB import PPBuilder, CaPPBuilder

from data import RDATA_PATH


def test_parse_pdb():
  PDB_FILE = Path(RDATA_PATH) / 'NMR_PDB' / 'NMR01' / '1CV9.pdb.gz'

  parser = PDBParser()
  with gzip.open(PDB_FILE, 'rt', encoding='utf-8') as fh:
    structure = parser.get_structure(PDB_FILE.stem[:-len('.pdb')], fh)

  for model in structure.get_models():
    for chain in model.get_chains():
      for residue in chain.get_residues():
        for atom in residue.get_atoms():
          print(atom)
          breakpoint()

  # Using C-N creterion
  ppb = PPBuilder()
  for pp in ppb.build_peptides(structure):
    print(pp.get_sequence())

  print('====')

  # Using CA-CA creterion
  ppb = CaPPBuilder()
  for pp in ppb.build_peptides(structure):
    print(pp.get_sequence())


if __name__ == '__main__':
  test_parse_pdb()
